from PIL import Image, ImageDraw, ImageFont
import os
from datetime import datetime
from PIL.ExifTags import TAGS, GPSTAGS
from typing import Tuple, Optional, Callable
import tkinter as tk
from tkinter import ttk, filedialog, messagebox

class ImageWatermarker:
    def __init__(self, 
                 date_format: str = '%Y-%m-%d %H:%M:%S',
                 text_color: Tuple[int, int, int] = (255, 165, 0),  # 橙色
                 font_size: int = 72,
                 font_path: Optional[str] = None,
                 margin: int = 40,
                 text_prefix: str = "拍摄时间: "):
        """
        初始化图片水印工具
        
        参数:
            date_format: 日期格式
            text_color: 文字颜色 RGB元组
            font_size: 字体大小
            font_path: 字体文件路径
            margin: 边距
            text_prefix: 前缀文本
        """
        self.date_format = date_format
        self.text_color = text_color
        self.font_size = font_size
        self.margin = margin
        self.text_prefix = text_prefix
        self.font = self._init_font(font_path, font_size)
        
        # 根据文字颜色自动计算描边颜色
        self.stroke_color = self._get_contrast_color(text_color)
        self.stroke_width = 2  # 固定描边宽度

    def _init_font(self, font_path: Optional[str], font_size: int) -> ImageFont.FreeTypeFont:
        """初始化字体"""
        try:
            if font_path:
                return ImageFont.truetype(font_path, font_size)
            
            # 首先尝试程序目录中的字体文件
            current_dir = os.path.dirname(os.path.abspath(__file__))
            font_files = [
                os.path.join(current_dir, "msyh.ttc"),
                os.path.join(current_dir, "simhei.ttf"),
                os.path.join(current_dir, "simsun.ttc")
            ]
            
            # 然后尝试系统字体
            system_font_files = ["msyh.ttc", "simhei.ttf", "simsun.ttc"]
            
            # 合并两个列表
            all_font_files = font_files + system_font_files
            
            for font_file in all_font_files:
                try:
                    return ImageFont.truetype(font_file, font_size)
                except:
                    continue
            
            print("警告: 未找到合适的中文字体，将使用默认字体")
            return ImageFont.load_default()
        except Exception as e:
            print(f"字体加载错误: {e}")
            return ImageFont.load_default()

    def _get_contrast_color(self, color: Tuple[int, int, int]) -> Tuple[int, int, int, int]:
        """根据文字颜色计算对比色（带透明度）"""
        # 计算颜色的亮度
        brightness = (color[0] * 299 + color[1] * 587 + color[2] * 114) / 1000
        
        # 设置透明度为40%（255 * 0.4 = 102）
        alpha = 102
        
        if brightness > 128:
            # 如果文字颜色较亮，使用暗色描边
            return (0, 0, 0, alpha)
        else:
            # 如果文字颜色较暗，使用亮色描边
            return (255, 255, 255, alpha)

    def get_capture_time(self, image_path: str) -> Optional[str]:
        """获取图片的拍摄时间"""
        try:
            with Image.open(image_path) as img:
                exif = img._getexif()
                if exif:
                    for tag_id, value in exif.items():
                        tag = TAGS.get(tag_id, tag_id)
                        if tag == 'DateTimeOriginal':
                            dt = datetime.strptime(value, '%Y:%m:%d %H:%M:%S')
                            return dt.strftime(self.date_format)
        except Exception as e:
            print(f"读取EXIF数据时出错 {image_path}: {e}")
        return None

    def add_timestamp(self, image_path: str, output_path: str) -> bool:
        try:
            with Image.open(image_path) as img:
                original_format = img.format
                original_size = img.size
                
                # 创建新图片，始终使用RGB模式
                new_img = Image.new('RGB', original_size, (255, 255, 255))
                new_img.paste(img, (0, 0))
                
                timestamp = self.get_capture_time(image_path)
                if timestamp:
                    text = f"{self.text_prefix}{timestamp}"
                    draw = ImageDraw.Draw(new_img)
                    
                    # 计算文本位置
                    bbox = draw.textbbox((0, 0), text, font=self.font)
                    text_width = bbox[2] - bbox[0]
                    text_height = bbox[3] - bbox[1]
                    
                    x = original_size[0] - text_width - self.margin
                    y = original_size[1] - text_height - self.margin
                    
                    # 绘制描边（带透明度）
                    stroke_color_rgb = self.stroke_color[:3]  # 获取RGB部分
                    stroke_alpha = self.stroke_color[3]       # 获取透明度
                    
                    # 计算带透明度的描边颜色
                    alpha_ratio = stroke_alpha / 255.0
                    stroke_color = tuple(int(c * alpha_ratio + 255 * (1 - alpha_ratio)) for c in stroke_color_rgb)
                    
                    # 绘制描边
                    for offset_x in range(-self.stroke_width, self.stroke_width + 1):
                        for offset_y in range(-self.stroke_width, self.stroke_width + 1):
                            if offset_x != 0 or offset_y != 0:
                                draw.text(
                                    (x + offset_x, y + offset_y),
                                    text,
                                    font=self.font,
                                    fill=stroke_color
                                )
                    
                    # 绘制主文本
                    draw.text((x, y), text, font=self.font, fill=self.text_color)
                
                # 保存图片
                new_img.save(
                    output_path,
                    'JPEG' if original_format in ['JPEG', 'JPG'] else original_format,
                    quality=100,
                    subsampling=0,
                    optimize=False,
                    exif=img.info.get('exif', b'')
                )
                
                print(f"已处理: {os.path.basename(image_path)}")
                return True
                
        except Exception as e:
            print(f"处理图片时出错 {image_path}: {e}")
            import traceback
            print(traceback.format_exc())
            return False

    def process_folder(self, input_folder: Optional[str] = None, 
                      output_folder: Optional[str] = None,
                      progress_callback: Optional[Callable[[float, str], None]] = None) -> None:
        """处理文件夹中的所有图片"""
        # 设置输入输出文件夹
        input_folder = input_folder or os.getcwd()
        if output_folder is None:
            output_folder = os.path.join(input_folder, "已添加时间戳")
        
        # 创建输出文件夹
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        
        # 获取需要处理的图片列表
        image_files = [f for f in os.listdir(input_folder) 
                      if f.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp'))]
        total_files = len(image_files)
        
        if total_files == 0:
            print("\n当前文件夹中没有找到图片文件。")
            return
        
        # 处理所有图片
        for i, filename in enumerate(image_files, 1):
            input_path = os.path.join(input_folder, filename)
            output_path = os.path.join(output_folder, filename)
            self.add_timestamp(input_path, output_path)
            
            # 更新进度并检查是否需要停止
            if progress_callback:
                progress = (i / total_files) * 100
                progress_callback(progress, f"正在处理: {filename} ({i}/{total_files})")
        
        print(f"\n处理完成！请在 '{output_folder}' 文件夹中查看处理后的图片。")

class WatermarkGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("图片时间戳工具")
        self.root.geometry("600x500")
        
        # 创建主框架
        main_frame = ttk.Frame(root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 日期格式选择
        ttk.Label(main_frame, text="日期格式:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.date_format_var = tk.StringVar()
        date_formats = [
            "完整格式 (2024-03-21 14:30:45)",
            "日期和时间 (2024-03-21 14:30)",
            "仅日期 (2024-03-21)",
            "中文格式 (2024年03月21日 14:30)",
            "中文日期 (2024年03月21日)"
        ]
        self.date_format_combo = ttk.Combobox(main_frame, textvariable=self.date_format_var, values=date_formats, width=30)
        self.date_format_combo.grid(row=0, column=1, sticky=tk.W, pady=5)
        self.date_format_combo.set(date_formats[0])
        
        # 文字颜色选择
        ttk.Label(main_frame, text="文字颜色:").grid(row=1, column=0, sticky=tk.W, pady=5)
        color_frame = ttk.Frame(main_frame)
        color_frame.grid(row=1, column=1, sticky=tk.W, pady=5)
        
        self.color_var = tk.StringVar()
        colors = ["橙色", "红色", "绿色", "蓝色", "白色", "黑色"]
        self.color_combo = ttk.Combobox(color_frame, textvariable=self.color_var, values=colors, width=20)
        self.color_combo.pack(side=tk.LEFT, padx=(0, 5))
        self.color_combo.set(colors[0])
        
        # 颜色预览框
        self.color_preview = tk.Canvas(color_frame, width=25, height=25, bg='orange')
        self.color_preview.pack(side=tk.LEFT, padx=5)
        
        # 自定义颜色按钮
        ttk.Button(color_frame, text="自定义颜色", command=self.choose_custom_color).pack(side=tk.LEFT, padx=5)
        
        # 绑定颜色选择事件
        self.color_combo.bind('<<ComboboxSelected>>', self.update_color_preview)
        
        # 字体大小
        ttk.Label(main_frame, text="字体大小:").grid(row=2, column=0, sticky=tk.W, pady=5)
        self.font_size_var = tk.StringVar(value="72")
        font_size_entry = ttk.Entry(main_frame, textvariable=self.font_size_var, width=10)
        font_size_entry.grid(row=2, column=1, sticky=tk.W, pady=5)
        
        # 边距设置
        ttk.Label(main_frame, text="边距(像素):").grid(row=3, column=0, sticky=tk.W, pady=5)
        self.margin_var = tk.StringVar(value="40")
        margin_entry = ttk.Entry(main_frame, textvariable=self.margin_var, width=10)
        margin_entry.grid(row=3, column=1, sticky=tk.W, pady=5)
        
        # 文字前缀
        ttk.Label(main_frame, text="文字前缀:").grid(row=4, column=0, sticky=tk.W, pady=5)
        self.prefix_var = tk.StringVar(value="拍摄时间：")
        prefix_entry = ttk.Entry(main_frame, textvariable=self.prefix_var, width=30)
        prefix_entry.grid(row=4, column=1, sticky=tk.W, pady=5)
        
        # 输入文件夹选择
        ttk.Label(main_frame, text="输入文件夹:").grid(row=5, column=0, sticky=tk.W, pady=5)
        self.input_folder_var = tk.StringVar(value=os.getcwd())
        input_folder_entry = ttk.Entry(main_frame, textvariable=self.input_folder_var, width=50)
        input_folder_entry.grid(row=5, column=1, sticky=tk.W, pady=5)
        ttk.Button(main_frame, text="浏览", command=self.browse_input_folder).grid(row=5, column=2, padx=5)
        
        # 输出文件夹选择
        ttk.Label(main_frame, text="输出文件夹:").grid(row=6, column=0, sticky=tk.W, pady=5)
        self.output_folder_var = tk.StringVar()
        output_folder_entry = ttk.Entry(main_frame, textvariable=self.output_folder_var, width=50)
        output_folder_entry.grid(row=6, column=1, sticky=tk.W, pady=5)
        ttk.Button(main_frame, text="浏览", command=self.browse_output_folder).grid(row=6, column=2, padx=5)
        
        # 进度条
        self.progress_var = tk.DoubleVar()
        self.progress_bar = ttk.Progressbar(main_frame, length=400, mode='determinate', variable=self.progress_var)
        self.progress_bar.grid(row=7, column=0, columnspan=3, pady=20)
        
        # 按钮框架
        button_frame = ttk.Frame(main_frame)
        button_frame.grid(row=8, column=0, columnspan=3, pady=10)
        
        # 开始处理按钮
        self.start_button = ttk.Button(button_frame, text="开始处理", command=self.process_images)
        self.start_button.pack(side=tk.LEFT, padx=5)
        
        # 停止处理按钮（初始状态为禁用）
        self.stop_button = ttk.Button(button_frame, text="停止处理", command=self.stop_processing, state='disabled')
        self.stop_button.pack(side=tk.LEFT, padx=5)
        
        # 状态标签
        self.status_var = tk.StringVar()
        self.status_label = ttk.Label(main_frame, textvariable=self.status_var)
        self.status_label.grid(row=9, column=0, columnspan=3)
        
        # 添加停止标志
        self.stop_flag = False

    def browse_input_folder(self):
        folder = filedialog.askdirectory(title="选择包含图片的文件夹")
        if folder:
            self.input_folder_var.set(folder)
            # 如果输出文件夹为空，自动设置默认输出路径
            if not self.output_folder_var.get():
                self.output_folder_var.set(os.path.join(folder, "已添加时间戳"))

    def browse_output_folder(self):
        folder = filedialog.askdirectory(title="选择输出文件夹")
        if folder:
            self.output_folder_var.set(folder)

    def get_color_tuple(self, color_name):
        """获取颜色RGB值"""
        default_colors = {
            "橙色": (255, 165, 0),
            "红色": (255, 0, 0),
            "绿色": (0, 255, 0),
            "蓝色": (0, 0, 255),
            "白色": (255, 255, 255),
            "黑色": (0, 0, 0)
        }
        
        # 检查是否是自定义颜色
        if hasattr(self, 'custom_colors') and color_name in self.custom_colors:
            return self.custom_colors[color_name]
            
        return default_colors.get(color_name, (255, 165, 0))

    def get_date_format(self, format_desc):
        format_map = {
            "完整格式 (2024-03-21 14:30:45)": '%Y-%m-%d %H:%M:%S',
            "日期和时间 (2024-03-21 14:30)": '%Y-%m-%d %H:%M',
            "仅日期 (2024-03-21)": '%Y-%m-%d',
            "中文格式 (2024年03月21日 14:30)": '%Y年%m月%d日 %H:%M',
            "中文日期 (2024年03月21日)": '%Y年%m月%d日'
        }
        return format_map.get(format_desc)

    def stop_processing(self):
        """停止处理图片"""
        self.stop_flag = True
        self.status_var.set("正在停止处理...")
        self.root.update()

    def process_images(self):
        try:
            # 重置停止标志
            self.stop_flag = False
            
            # 禁用开始按钮，启用停止按钮
            self.start_button.configure(state='disabled')
            self.stop_button.configure(state='normal')
            self.root.update()
            
            settings = {
                'date_format': self.get_date_format(self.date_format_var.get()),
                'text_color': self.get_color_tuple(self.color_var.get()),
                'font_size': int(self.font_size_var.get()),
                'margin': int(self.margin_var.get()),
                'text_prefix': self.prefix_var.get(),
                'input_folder': self.input_folder_var.get().strip(),
                'output_folder': self.output_folder_var.get().strip()
            }
            
            # 检查输入文件夹
            if not settings['input_folder']:
                messagebox.showerror("错误", "请选择输入文件夹！")
                return
            
            if not os.path.exists(settings['input_folder']):
                messagebox.showerror("错误", f"输入文件夹不存在：\n{settings['input_folder']}")
                return
            
            # 设置输出文件夹
            if not settings['output_folder']:
                settings['output_folder'] = os.path.join(settings['input_folder'], "已添加时间戳")
                self.output_folder_var.set(settings['output_folder'])
            
            # 确保输出文件夹的父目录存在
            output_parent = os.path.dirname(settings['output_folder'])
            if output_parent and not os.path.exists(output_parent):
                try:
                    os.makedirs(output_parent)
                except Exception as e:
                    messagebox.showerror("错误", f"无法创建输出路径：\n{output_parent}\n\n错误信息：{str(e)}")
                    return
            
            # 创建输出文件夹
            try:
                os.makedirs(settings['output_folder'], exist_ok=True)
            except Exception as e:
                messagebox.showerror("错误", f"无法创建输出文件夹：\n{settings['output_folder']}\n\n错误信息：{str(e)}")
                return
            
            # 检查输入文件夹中是否有图片
            image_files = [f for f in os.listdir(settings['input_folder']) 
                          if f.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp'))]
            if not image_files:
                messagebox.showwarning("警告", f"在输入文件夹中没有找到图片文件：\n{settings['input_folder']}")
                return
            
            # 创建水印工具实例
            watermarker = ImageWatermarker(
                date_format=settings['date_format'],
                text_color=settings['text_color'],
                font_size=settings['font_size'],
                margin=settings['margin'],
                text_prefix=settings['text_prefix']
            )
            
            # 修改进度回调函数
            def update_progress(progress, status):
                if self.stop_flag:
                    raise InterruptedError("用户停止了处理")
                self.progress_var.set(progress)
                self.status_var.set(status)
                self.root.update()
            
            # 处理图片
            watermarker.process_folder(
                input_folder=settings['input_folder'],
                output_folder=settings['output_folder'],
                progress_callback=update_progress
            )
            
            # 完成处理
            self.progress_var.set(100)
            self.status_var.set("处理完成！")
            messagebox.showinfo("完成", f"所有图片处理完成！\n\n输出文件夹：\n{settings['output_folder']}")
            
        except InterruptedError:
            self.status_var.set("处理已停止")
            messagebox.showinfo("已停止", "图片处理已被用户停止")
        except Exception as e:
            import traceback
            error_msg = f"处理过程中出现错误：\n{str(e)}\n\n详细信息：\n{traceback.format_exc()}"
            messagebox.showerror("错误", error_msg)
        finally:
            try:
                # 恢复按钮状态
                if self.root.winfo_exists():  # 检查窗口是否还存在
                    self.start_button.configure(state='normal')
                    self.stop_button.configure(state='disabled')
                    self.root.update()
            except:
                pass  # 忽略任何恢复按钮状态时的错误

    def update_color_preview(self, event=None):
        """更新颜色预览"""
        color = self.get_color_tuple(self.color_var.get())
        # 将RGB颜色转换为十六进制格式
        hex_color = '#{:02x}{:02x}{:02x}'.format(*color)
        self.color_preview.configure(bg=hex_color)

    def choose_custom_color(self):
        """打开颜色选择器"""
        from tkinter import colorchooser
        color = colorchooser.askcolor(title="选择自定义颜色")
        if color[0]:  # color[0]是RGB值，color[1]是十六进制值
            # 添加自定义颜色到下拉列表
            rgb_color = tuple(map(int, color[0]))
            custom_name = f"自定义 ({rgb_color[0]},{rgb_color[1]},{rgb_color[2]})"
            
            # 更新颜色映射
            self.custom_colors = getattr(self, 'custom_colors', {})
            self.custom_colors[custom_name] = rgb_color
            
            # 更新下拉列表
            current_values = list(self.color_combo['values'])
            if custom_name not in current_values:
                current_values.append(custom_name)
                self.color_combo['values'] = current_values
            
            # 选择新颜色
            self.color_combo.set(custom_name)
            self.update_color_preview()

if __name__ == "__main__":
    root = tk.Tk()
    app = WatermarkGUI(root)
    root.mainloop()
